{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tensorflow-compress.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyM70KjgiH+IuhSm+julxt5Q"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7RWDC4Q5RsXC"
      },
      "source": [
        "# tensorflow-compress"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kk-R-bErR7gS"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/byronknoll/tensorflow-compress/blob/master/tensorflow-compress.ipynb)\n",
        "\n",
        "Made by Byron Knoll. GitHub repository: https://github.com/byronknoll/tensorflow-compress\n",
        "\n",
        "### Description\n",
        "\n",
        "tensorflow-compress performs lossless data compression using neural networks in TensorFlow. It can run on GPUs with a large batch size to get a substantial speed improvement. It is made using Colab, which should make it easy to run through a web browser. You can choose a file, perform compression (or decompression), and download the result.\n",
        "\n",
        "tensorflow-compress is open source and the code should hopefully be easy to understand and modify. Feel free to experiment with the code and create pull requests with improvements.\n",
        "\n",
        "The neural network is trained from scratch during compression and decompression, so the model weights do not need to be stored. Arithmetic coding is used to encode the model predictions to a file.\n",
        "\n",
        "Feel free to contact me at byron@byronknoll.com if you have any questions.\n",
        "\n",
        "### Instructions:\n",
        "\n",
        "Basic usage: configure all the fields in the \"Parameters\" section and select Runtime->Run All.\n",
        "\n",
        "Advanced usage: save a copy of this notebook and modify the code.\n",
        "\n",
        "### Related Projects\n",
        "*   [NNCP](https://bellard.org/nncp/) - this uses a similar LSTM architecture to tensorflow-compress. It is limited to running only on CPUs.\n",
        "*   [lstm-compress](https://github.com/byronknoll/lstm-compress) - similar to NNCP, but has a batch size limit of one (so it is significantly slower).\n",
        "*   [cmix](http://www.byronknoll.com/cmix.html) - shares the same LSTM code as lstm-compress, but contains a bunch of other components to get better compression rate.\n",
        "*   [DeepZip](https://github.com/mohit1997/DeepZip) - this also performs compression using TensorFlow. However, it has some substantial architecture differences to tensorflow-compress: it uses pretraining (using multiple passes over the training data) and stores the model weights in the compressed file.\n",
        "\n",
        "### Benchmarks\n",
        "These benchmarks were performed using tensorflow-compress v3 with the default parameter settings. Some parameters differ between enwik8 and enwik9 as noted in the parameter comments. Colab Pro was used with Tesla V100 GPU. Compression time and decompression time are approximately the same.\n",
        "*   enwik8: compressed to 16,128,954 bytes in 32,113.38 seconds. NNCP preprocessing time: 206.38 seconds. Dictionary size: 65,987 bytes.\n",
        "*   enwik9: compressed to 118,938,744 bytes in 297,505.98 seconds. NNCP preprocessing time: 2,598.77 seconds. Dictionary size: 79,876 bytes. Since Colab has a 24 hour time limit, the preprocessed enwik9 file was split into four parts using [this notebook](https://colab.sandbox.google.com/github/byronknoll/tensorflow-compress/blob/master/nncp-splitter.ipynb). The \"checkpoint\" option was used to save/load model weights between processing each part. For the first part, start_learning_rate=0.0007 and end_learning_rate=0.0005 was used. For the remaining three parts, a constant 0.00035 learning rate was used.\n",
        "\n",
        "See the [Large Text Compression Benchmark](http://mattmahoney.net/dc/text.html) for more information about the test files and a comparison with other programs.\n",
        "\n",
        "### Versions\n",
        "* v3 - released November 28, 2020. Changes from v2:\n",
        "  * Parameter tuning\n",
        "  * [New notebook](https://colab.sandbox.google.com/github/byronknoll/tensorflow-compress/blob/master/nncp-splitter.ipynb) for file splitting\n",
        "  * Support for learning rate decay\n",
        "* v2 - released September 6, 2020. Changes from v1:\n",
        "  * 16 bit floats for improved speed\n",
        "  * Weight updates occur at every timestep (instead of at spaced intervals)\n",
        "  * Support for saving/loading model weights\n",
        "* v1 - released July 20, 2020."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xQ5gF5mLs0Ew"
      },
      "source": [
        "## Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZUc1x3vIs4hF",
        "cellView": "form"
      },
      "source": [
        "batch_size = 96 #@param {type:\"integer\"}\n",
        "#@markdown >_This will split the file into N batches, and process them in parallel. Increasing this will improve speed but can make compression rate worse. Make this a multiple of 8 to improve speed on certain GPUs._\n",
        "seq_length = 11 #@param {type:\"integer\"}\n",
        "#@markdown >_This determines the horizon for back propagation through time. Reducing this will improve speed, but can make compression rate worse._\n",
        "rnn_units =  1400 #@param {type:\"integer\"}\n",
        "#@markdown >_This is the number of units to use within each LSTM layer. Reducing this will improve speed, but can make compression rate worse. Make this a multiple of 8 to improve speed on certain GPUs._\n",
        "num_layers = 6 #@param {type:\"integer\"}\n",
        "#@markdown >_This is the number of LSTM layers to use. Reducing this will improve speed, but can make compression rate worse._\n",
        "start_learning_rate = 0.00075 #@param {type:\"number\"}\n",
        "#@markdown >_Learning rate for Adam optimizer. Recommended value for enwik8: 0.00075. For enwik9, see the notes in the \"Benchmarks\" section for the recommended learning rate._\n",
        "end_learning_rate = 0.00075 #@param {type:\"number\"}\n",
        "#@markdown >_Typically this should be set to the same value as the \"start_learning_rate\" parameter above. If this is set to a different value, the learning rate will start at \"start_learning_rate\" and linearly change to \"end_learning_rate\" by the end of the file. For large files this could be useful for learning rate decay._\n",
        "mode = 'compress' #@param [\"compress\", \"decompress\", \"both\", \"preprocess_only\"]\n",
        "#@markdown >_Whether to run compression only, decompression only, or both. \"preprocess_only\" will only run preprocessing and skip compression._\n",
        "preprocess = 'nncp' #@param [\"cmix\", \"nncp\", \"nncp-done\", \"none\"]\n",
        "#@markdown >_The choice of preprocessor. NNCP works better on enwik8/enwik9. NNCP preprocessing is slower since it constructs a custom dictionary, while cmix uses a pretrained dictionary. \"nncp_done\" is used for files which have already been preprocessed by NNCP (the dictionary must also be included)._\n",
        "n_words = 8192 #@param {type:\"integer\"}\n",
        "#@markdown >_Only used for NNCP preprocessor: this is the approximative maximum number of words of the dictionary. Recommended value for enwik8/enwik9: 8192._\n",
        "min_freq = 64 #@param {type:\"integer\"}\n",
        "#@markdown >_Only used for NNCP preprocessor: this is the minimum frequency of the selected words. Recommended value for enwik8: 64, enwik9: 512._\n",
        "path_to_file = \"enwik8\" #@param [\"enwik4\", \"enwik6\", \"enwik8\", \"enwik9\", \"custom\"]\n",
        "#@markdown >_Name of the file to compress or decompress. If \"custom\" is selected, use the next parameter to set a custom path._\n",
        "custom_path = '' #@param {type:\"string\"}\n",
        "#@markdown >_Use this if the previous parameter was set to \"custom\". Set this to the name of the file you want to compress/decompress. You can transfer files using the \"http_path\" or \"local_upload\" options below._\n",
        "http_path = '' #@param {type:\"string\"}\n",
        "#@markdown >_The file from this URL will be downloaded. It is recommended to use Google Drive URLs to get fast transfer speed. Use this format for Google Drive files: https://drive.google.com/uc?id= and paste the file ID at the end of the URL. You can find the file ID from the \"Get Link\" URL in Google Drive. You can enter multiple URLs here, space separated._\n",
        "local_upload = False #@param {type:\"boolean\"}\n",
        "#@markdown >_If enabled, you will be prompted in the \"Setup Files\" section to select files to upload from your local computer. You can upload multiple files. Note: the upload speed can be quite slow (use \"http_path\" for better transfer speeds)._\n",
        "download_option = \"no_download\" #@param [\"no_download\", \"local\", \"google_drive\"]\n",
        "#@markdown >_If this is set to \"local\", the output files will be downloaded to your computer after compression/decompression. If set to \"google_drive\", they will be copied to your Google Drive account (which is significantly faster than downloading locally)._\n",
        "checkpoint = False #@param {type:\"boolean\"}\n",
        "#@markdown >_If this is enabled, a checkpoint of the model weights will be downloaded (using the \"download_option\" parameter). This can be useful for getting around session time limits for Colab, by splitting files into multiple segments and saving/loading the model weights between each segment. Checkpoints (if present) will automatically be loaded when starting compression._\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yezcaYiWzhoV"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NtRgaznIzneg",
        "cellView": "form"
      },
      "source": [
        "#@title Imports\n",
        "\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import random\n",
        "from google.colab import files\n",
        "import time\n",
        "import math\n",
        "import sys\n",
        "import subprocess\n",
        "import contextlib\n",
        "import os\n",
        "import tensorflow_addons as tfa\n",
        "from tensorflow.keras.mixed_precision import experimental as mixed_precision\n",
        "from google.colab import drive\n",
        "os.environ['TF_DETERMINISTIC_OPS'] = '1'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2795ZUv4BPQi",
        "cellView": "form"
      },
      "source": [
        "#@title System Info\n",
        "\n",
        "def system_info():\n",
        "  \"\"\"Prints out system information.\"\"\"\n",
        "  gpu_info = !nvidia-smi\n",
        "  gpu_info = '\\n'.join(gpu_info)\n",
        "  if gpu_info.find('failed') >= 0:\n",
        "    print('Select the Runtime → \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "    print('and then re-execute this cell.')\n",
        "  else:\n",
        "    print(gpu_info)\n",
        "  print(\"TensorFlow version: \", tf.__version__)\n",
        "  !lscpu |grep 'Model name'\n",
        "  !cat /proc/meminfo | head -n 3\n",
        "\n",
        "system_info()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8TjO6DWgoA7i",
        "cellView": "form"
      },
      "source": [
        "#@title Mount Google Drive\n",
        "if download_option == \"google_drive\":\n",
        "  drive.mount('/content/gdrive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8FL-Z7yMU4FF",
        "cellView": "form"
      },
      "source": [
        "#@title Setup Files\n",
        "\n",
        "!mkdir -p \"data\"\n",
        "\n",
        "if local_upload:\n",
        "  %cd data\n",
        "  files.upload()\n",
        "  %cd ..\n",
        "\n",
        "if path_to_file == 'enwik8' or path_to_file == 'enwik6' or path_to_file == 'enwik4':\n",
        "  %cd data\n",
        "  !gdown --id 1BUbuEUhPOBaVZDdOh0KG8hxvIDgsyiZp\n",
        "  !unzip enwik8.zip\n",
        "  !head -c 1000000 enwik8 > enwik6\n",
        "  !head -c 10000 enwik8 > enwik4\n",
        "  path_to_file = 'data/' + path_to_file\n",
        "  %cd ..\n",
        "\n",
        "if path_to_file == 'enwik9':\n",
        "  %cd data\n",
        "  !gdown --id 1D2gCmf9AlXIBP62ARhy0XcIuIolOTRAE\n",
        "  !unzip enwik9.zip\n",
        "  path_to_file = 'data/' + path_to_file\n",
        "  %cd ..\n",
        "\n",
        "if path_to_file == 'custom':\n",
        "  path_to_file = 'data/' + custom_path\n",
        "\n",
        "if http_path:\n",
        "  %cd data\n",
        "  paths = http_path.split()\n",
        "  for path in paths:\n",
        "    !gdown $path\n",
        "  %cd ..\n",
        "\n",
        "if preprocess == 'cmix':\n",
        "  !gdown --id 1qa7K28tlUDs9GGYbaL_iE9M4m0L1bYm9\n",
        "  !unzip cmix-v18.zip\n",
        "  %cd cmix\n",
        "  !make\n",
        "  %cd ..\n",
        "\n",
        "if preprocess == 'nncp' or preprocess == 'nncp-done':\n",
        "  !gdown --id 1EzVPbRkBIIbgOzvEMeM0YpibDi2R4SHD\n",
        "  !tar -xf nncp-2019-11-16.tar.gz\n",
        "  %cd nncp-2019-11-16/\n",
        "  !make preprocess\n",
        "  %cd .."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f5ZRPnh-2dCz",
        "cellView": "form"
      },
      "source": [
        "#@title Model Architecture\n",
        "\n",
        "def build_model(vocab_size):\n",
        "  \"\"\"Builds the model architecture.\n",
        "\n",
        "    Args:\n",
        "      vocab_size: Int, size of the vocabulary.\n",
        "  \"\"\"\n",
        "  policy = mixed_precision.Policy('mixed_float16')\n",
        "  mixed_precision.set_policy(policy)\n",
        "  inputs = [\n",
        "    tf.keras.Input(batch_input_shape=[batch_size, seq_length, vocab_size])]\n",
        "  # In addition to the primary input, there are also two \"state\" inputs for each\n",
        "  # layer of the network.\n",
        "  for i in range(num_layers):\n",
        "    inputs.append(tf.keras.Input(shape=(None,)))\n",
        "    inputs.append(tf.keras.Input(shape=(None,)))\n",
        "  # Skip connections will be used to connect each LSTM layer output to the final\n",
        "  # output layer. Each LSTM layer will get as input both the original input and\n",
        "  # the output of the previous layer.\n",
        "  skip_connections = []\n",
        "  # In addition to the softmax output, there are also two \"state\" outputs for\n",
        "  # each layer of the network.\n",
        "  outputs = []\n",
        "  predictions, state_h, state_c = tf.keras.layers.LSTM(rnn_units,\n",
        "                          return_sequences=True,\n",
        "                          return_state=True,\n",
        "                          recurrent_initializer='glorot_uniform',\n",
        "                          )(inputs[0], initial_state=[\n",
        "                          tf.cast(inputs[1], tf.float16),\n",
        "                          tf.cast(inputs[2], tf.float16)])\n",
        "  skip_connections.append(predictions)\n",
        "  outputs.append(state_h)\n",
        "  outputs.append(state_c)\n",
        "  for i in range(num_layers - 1):\n",
        "    layer_input = tf.keras.layers.concatenate(\n",
        "        [inputs[0], skip_connections[-1]])\n",
        "    predictions, state_h, state_c = tf.keras.layers.LSTM(rnn_units,\n",
        "        return_sequences=True,\n",
        "        return_state=True,\n",
        "        recurrent_initializer='glorot_uniform')(\n",
        "          layer_input, initial_state=[tf.cast(inputs[i*2+3], tf.float16),\n",
        "                                      tf.cast(inputs[i*2+4], tf.float16)])\n",
        "    skip_connections.append(predictions)\n",
        "    outputs.append(state_h)\n",
        "    outputs.append(state_c)\n",
        "  # The dense output layer only needs to be computed for the last timestep, so\n",
        "  # we can discard the earlier outputs.\n",
        "  last_timestep = []\n",
        "  for i in range(num_layers):\n",
        "    last_timestep.append(tf.slice(skip_connections[i], [0, seq_length - 1, 0],\n",
        "                                [batch_size, 1, rnn_units]))\n",
        "  if num_layers == 1:\n",
        "    layer_input = last_timestep[0]\n",
        "  else:\n",
        "    layer_input = tf.keras.layers.concatenate(last_timestep)\n",
        "  dense = tf.keras.layers.Dense(vocab_size, name='dense_logits')(layer_input)\n",
        "  output = tf.keras.layers.Activation('softmax', dtype='float32',\n",
        "                                      name='predictions')(dense)\n",
        "  outputs.insert(0, output)\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NlytnNHdfbIM",
        "cellView": "form"
      },
      "source": [
        "#@title Compression Library\n",
        "\n",
        "def get_symbol(index, length, freq, coder, compress, data):\n",
        "  \"\"\"Runs arithmetic coding and returns the next symbol.\n",
        "\n",
        "  Args:\n",
        "    index: Int, position of the symbol in the file.\n",
        "    length: Int, size limit of the file.\n",
        "    freq: ndarray, predicted symbol probabilities.\n",
        "    coder: this is the arithmetic coder.\n",
        "    compress: Boolean, True if compressing, False if decompressing.\n",
        "    data: List containing each symbol in the file.\n",
        "  \n",
        "  Returns:\n",
        "    The next symbol, or 0 if \"index\" is over the file size limit.\n",
        "  \"\"\"\n",
        "  symbol = 0\n",
        "  if index < length:\n",
        "    if compress:\n",
        "      symbol = data[index]\n",
        "      coder.write(freq, symbol)\n",
        "    else:\n",
        "      symbol = coder.read(freq)\n",
        "      data[index] = symbol\n",
        "  return symbol\n",
        "\n",
        "def train(pos, seq_input, length, vocab_size, coder, model, optimizer, compress,\n",
        "          data, states):\n",
        "  \"\"\"Runs one training step.\n",
        "\n",
        "  Args:\n",
        "    pos: Int, position in the file for the current symbol for the *first* batch.\n",
        "    seq_input: Tensor, containing the last seq_length inputs for the model.\n",
        "    length: Int, size limit of the file.\n",
        "    vocab_size: Int, size of the vocabulary.\n",
        "    coder: this is the arithmetic coder.\n",
        "    model: the model to generate predictions.\n",
        "    optimizer: optimizer used to train the model.\n",
        "    compress: Boolean, True if compressing, False if decompressing.\n",
        "    data: List containing each symbol in the file.\n",
        "    states: List containing state information for the layers of the model.\n",
        "  \n",
        "  Returns:\n",
        "    seq_input: Tensor, containing the last seq_length inputs for the model.\n",
        "    cross_entropy: cross entropy numerator.\n",
        "    denom: cross entropy denominator.\n",
        "  \"\"\"\n",
        "  loss = cross_entropy = denom = 0\n",
        "  split = math.ceil(length / batch_size)\n",
        "  # Keep track of operations while running the forward pass for automatic\n",
        "  # differentiation.\n",
        "  with tf.GradientTape() as tape:\n",
        "    # The model inputs contain both seq_input and the states for each layer.\n",
        "    inputs = states.pop(0)\n",
        "    inputs.insert(0, seq_input)\n",
        "    # Run the model (for all batches in parallel) to get predictions for the\n",
        "    # next characters.\n",
        "    outputs = model(inputs)\n",
        "    predictions = outputs.pop(0)\n",
        "    states.append(outputs)\n",
        "    p = predictions.numpy()\n",
        "    symbols = []\n",
        "    # When the last batch reaches the end of the file, we start giving it \"0\"\n",
        "    # as input. We use a mask to prevent this from influencing the gradients.\n",
        "    mask = []\n",
        "    # Go over each batch to run the arithmetic coding and prepare the next\n",
        "    # input.\n",
        "    for i in range(batch_size):\n",
        "      # The \"10000000\" is used to convert floats into large integers (since\n",
        "      # the arithmetic coder works on integers).\n",
        "      freq = np.cumsum(p[i][0] * 10000000 + 1)\n",
        "      index = pos + 1 + i * split\n",
        "      symbol = get_symbol(index, length, freq, coder, compress, data)\n",
        "      symbols.append(symbol)\n",
        "      if index < length:\n",
        "        prob = p[i][0][symbol]\n",
        "        if prob <= 0:\n",
        "          # Set a small value to avoid error with log2.\n",
        "          prob = 0.000001\n",
        "        cross_entropy += math.log2(prob)\n",
        "        denom += 1\n",
        "        mask.append(1.0)\n",
        "      else:\n",
        "        mask.append(0.0)\n",
        "    # \"input_one_hot\" will be used both for the loss function and for the next\n",
        "    # input.\n",
        "    input_one_hot = tf.expand_dims(tf.one_hot(symbols, vocab_size), 1)\n",
        "    loss = tf.keras.losses.categorical_crossentropy(\n",
        "        input_one_hot, predictions, from_logits=False) * tf.expand_dims(\n",
        "            tf.convert_to_tensor(mask), 1)\n",
        "    # Remove the oldest input and append the new one.\n",
        "    seq_input = tf.slice(seq_input, [0, 1, 0],\n",
        "                         [batch_size, seq_length - 1, vocab_size])\n",
        "    seq_input = tf.concat([seq_input, input_one_hot], 1)\n",
        "  # Run the backwards pass to update model weights.\n",
        "  grads = tape.gradient(loss, model.trainable_variables)\n",
        "  # Gradient clipping to make training more robust.\n",
        "  capped_grads = [tf.clip_by_value(grad, -5., 5.) for grad in grads]\n",
        "  optimizer.apply_gradients(zip(capped_grads, model.trainable_variables))\n",
        "  return (seq_input, cross_entropy, denom)\n",
        "\n",
        "def reset_seed():\n",
        "  \"\"\"Initializes various random seeds to help with determinism.\"\"\"\n",
        "  SEED = 1234\n",
        "  os.environ['PYTHONHASHSEED']=str(SEED)\n",
        "  random.seed(SEED)\n",
        "  np.random.seed(SEED)\n",
        "  tf.random.set_seed(SEED)\n",
        "\n",
        "def download(path):\n",
        "  \"\"\"Downloads the file at the specified path.\"\"\"\n",
        "  if download_option == 'local':\n",
        "    files.download(path)\n",
        "  elif download_option == 'google_drive':\n",
        "    !cp -f $path /content/gdrive/My\\ Drive\n",
        "\n",
        "def process(compress, length, vocab_size, coder, data):\n",
        "  \"\"\"This runs compression/decompression.\n",
        "\n",
        "  Args:\n",
        "    compress: Boolean, True if compressing, False if decompressing.\n",
        "    length: Int, size limit of the file.\n",
        "    vocab_size: Int, size of the vocabulary.\n",
        "    coder: this is the arithmetic coder.\n",
        "    data: List containing each symbol in the file.\n",
        "  \"\"\"\n",
        "  start = time.time()\n",
        "  reset_seed()\n",
        "  model = build_model(vocab_size = vocab_size)\n",
        "  checkpoint_path = tf.train.latest_checkpoint('./data')\n",
        "  if checkpoint_path:\n",
        "    model.load_weights(checkpoint_path)\n",
        "  model.summary()\n",
        "\n",
        "  # Try to split the file into equal size pieces for the different batches. The\n",
        "  # last batch may have fewer characters if the file can't be split equally.\n",
        "  split = math.ceil(length / batch_size)\n",
        "\n",
        "  learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(\n",
        "      start_learning_rate,\n",
        "      split,\n",
        "      end_learning_rate,\n",
        "      power=1.0)\n",
        "  optimizer = tf.keras.optimizers.Adam(\n",
        "      learning_rate=learning_rate_fn, beta_1=0, beta_2=0.9999, epsilon=1e-5)\n",
        "\n",
        "  hidden = model.reset_states()\n",
        "  # Use a uniform distribution for predicting the first batch of symbols. The\n",
        "  # \"10000000\" is used to convert floats into large integers (since the\n",
        "  # arithmetic coder works on integers).\n",
        "  freq = np.cumsum(np.full(vocab_size, (1.0 / vocab_size)) * 10000000 + 1)\n",
        "  # Construct the first set of input characters for training.\n",
        "  symbols = []\n",
        "  for i in range(batch_size):\n",
        "    symbols.append(get_symbol(i*split, length, freq, coder, compress, data))\n",
        "  input_one_hot = tf.expand_dims(tf.one_hot(symbols, vocab_size), 1)\n",
        "  # Replicate the input tensor seq_length times, to match the input format.\n",
        "  seq_input = tf.tile(input_one_hot, [1, seq_length, 1])\n",
        "  pos = cross_entropy = denom = last_output = 0\n",
        "  template = '{:0.2f}%\\tcross entropy: {:0.2f}\\ttime: {:0.2f}'\n",
        "  # This will keep track of layer states. Initialize them to zeros.\n",
        "  states = []\n",
        "  for i in range(seq_length):\n",
        "    states.append([tf.zeros([batch_size, rnn_units])] * (num_layers * 2))\n",
        "  # Keep repeating the training step until we get to the end of the file.\n",
        "  while pos < split:\n",
        "    seq_input, ce, d = train(pos, seq_input, length, vocab_size, coder, model,\n",
        "                             optimizer, compress, data, states)\n",
        "    cross_entropy += ce\n",
        "    denom += d\n",
        "    pos += 1\n",
        "    time_diff = time.time() - start\n",
        "    # If it has been over 20 seconds since the last status message, display a\n",
        "    # new one.\n",
        "    if time_diff - last_output > 20:\n",
        "      last_output = time_diff\n",
        "      percentage = 100 * pos / split\n",
        "      if percentage >= 100: continue\n",
        "      print(template.format(percentage, -cross_entropy / denom, time_diff))\n",
        "  if compress:\n",
        "    coder.finish()\n",
        "  print(template.format(100, -cross_entropy / length, time.time() - start))\n",
        "  system_info()\n",
        "  if mode != \"both\" or not compress:\n",
        "    model.save_weights('./data/model')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6WEbfVSEKilf",
        "cellView": "form"
      },
      "source": [
        "#@title Arithmetic Coding Library\n",
        "\n",
        "# \n",
        "# Reference arithmetic coding\n",
        "# Copyright (c) Project Nayuki\n",
        "# \n",
        "# https://www.nayuki.io/page/reference-arithmetic-coding\n",
        "# https://github.com/nayuki/Reference-arithmetic-coding\n",
        "# \n",
        "\n",
        "import sys\n",
        "python3 = sys.version_info.major >= 3\n",
        "\n",
        "\n",
        "# ---- Arithmetic coding core classes ----\n",
        "\n",
        "# Provides the state and behaviors that arithmetic coding encoders and decoders share.\n",
        "class ArithmeticCoderBase(object):\n",
        "\t\n",
        "\t# Constructs an arithmetic coder, which initializes the code range.\n",
        "\tdef __init__(self, numbits):\n",
        "\t\tif numbits < 1:\n",
        "\t\t\traise ValueError(\"State size out of range\")\n",
        "\t\t\n",
        "\t\t# -- Configuration fields --\n",
        "\t\t# Number of bits for the 'low' and 'high' state variables. Must be at least 1.\n",
        "\t\t# - Larger values are generally better - they allow a larger maximum frequency total (maximum_total),\n",
        "\t\t#   and they reduce the approximation error inherent in adapting fractions to integers;\n",
        "\t\t#   both effects reduce the data encoding loss and asymptotically approach the efficiency\n",
        "\t\t#   of arithmetic coding using exact fractions.\n",
        "\t\t# - But larger state sizes increase the computation time for integer arithmetic,\n",
        "\t\t#   and compression gains beyond ~30 bits essentially zero in real-world applications.\n",
        "\t\t# - Python has native bigint arithmetic, so there is no upper limit to the state size.\n",
        "\t\t#   For Java and C++ where using native machine-sized integers makes the most sense,\n",
        "\t\t#   they have a recommended value of num_state_bits=32 as the most versatile setting.\n",
        "\t\tself.num_state_bits = numbits\n",
        "\t\t# Maximum range (high+1-low) during coding (trivial), which is 2^num_state_bits = 1000...000.\n",
        "\t\tself.full_range = 1 << self.num_state_bits\n",
        "\t\t# The top bit at width num_state_bits, which is 0100...000.\n",
        "\t\tself.half_range = self.full_range >> 1  # Non-zero\n",
        "\t\t# The second highest bit at width num_state_bits, which is 0010...000. This is zero when num_state_bits=1.\n",
        "\t\tself.quarter_range = self.half_range >> 1  # Can be zero\n",
        "\t\t# Minimum range (high+1-low) during coding (non-trivial), which is 0010...010.\n",
        "\t\tself.minimum_range = self.quarter_range + 2  # At least 2\n",
        "\t\t# Maximum allowed total from a frequency table at all times during coding. This differs from Java\n",
        "\t\t# and C++ because Python's native bigint avoids constraining the size of intermediate computations.\n",
        "\t\tself.maximum_total = self.minimum_range\n",
        "\t\t# Bit mask of num_state_bits ones, which is 0111...111.\n",
        "\t\tself.state_mask = self.full_range - 1\n",
        "\t\t\n",
        "\t\t# -- State fields --\n",
        "\t\t# Low end of this arithmetic coder's current range. Conceptually has an infinite number of trailing 0s.\n",
        "\t\tself.low = 0\n",
        "\t\t# High end of this arithmetic coder's current range. Conceptually has an infinite number of trailing 1s.\n",
        "\t\tself.high = self.state_mask\n",
        "\t\n",
        "\t\n",
        "\t# Updates the code range (low and high) of this arithmetic coder as a result\n",
        "\t# of processing the given symbol with the given frequency table.\n",
        "\t# Invariants that are true before and after encoding/decoding each symbol\n",
        "\t# (letting full_range = 2^num_state_bits):\n",
        "\t# - 0 <= low <= code <= high < full_range. ('code' exists only in the decoder.)\n",
        "\t#   Therefore these variables are unsigned integers of num_state_bits bits.\n",
        "\t# - low < 1/2 * full_range <= high.\n",
        "\t#   In other words, they are in different halves of the full range.\n",
        "\t# - (low < 1/4 * full_range) || (high >= 3/4 * full_range).\n",
        "\t#   In other words, they are not both in the middle two quarters.\n",
        "\t# - Let range = high - low + 1, then full_range/4 < minimum_range\n",
        "\t#   <= range <= full_range. These invariants for 'range' essentially\n",
        "\t#   dictate the maximum total that the incoming frequency table can have.\n",
        "\tdef update(self, freqs, symbol):\n",
        "\t\t# State check\n",
        "\t\tlow = self.low\n",
        "\t\thigh = self.high\n",
        "\t\t# if low >= high or (low & self.state_mask) != low or (high & self.state_mask) != high:\n",
        "\t\t# \traise AssertionError(\"Low or high out of range\")\n",
        "\t\trange = high - low + 1\n",
        "\t\t# if not (self.minimum_range <= range <= self.full_range):\n",
        "\t\t# \traise AssertionError(\"Range out of range\")\n",
        "\t\t\n",
        "\t\t# Frequency table values check\n",
        "\t\ttotal = int(freqs[-1])\n",
        "\t\tsymlow = int(freqs[symbol-1]) if symbol > 0 else 0\n",
        "\t\tsymhigh = int(freqs[symbol])\n",
        "\t\t#total = freqs.get_total()\n",
        "\t\t#symlow = freqs.get_low(symbol)\n",
        "\t\t#symhigh = freqs.get_high(symbol)\n",
        "\t\t# if symlow == symhigh:\n",
        "\t\t# \traise ValueError(\"Symbol has zero frequency\")\n",
        "\t\t# if total > self.maximum_total:\n",
        "\t\t# \traise ValueError(\"Cannot code symbol because total is too large\")\n",
        "\t\t\n",
        "\t\t# Update range\n",
        "\t\tnewlow  = low + symlow  * range // total\n",
        "\t\tnewhigh = low + symhigh * range // total - 1\n",
        "\t\tself.low = newlow\n",
        "\t\tself.high = newhigh\n",
        "\t\t\n",
        "\t\t# While low and high have the same top bit value, shift them out\n",
        "\t\twhile ((self.low ^ self.high) & self.half_range) == 0:\n",
        "\t\t\tself.shift()\n",
        "\t\t\tself.low  = ((self.low  << 1) & self.state_mask)\n",
        "\t\t\tself.high = ((self.high << 1) & self.state_mask) | 1\n",
        "\t\t# Now low's top bit must be 0 and high's top bit must be 1\n",
        "\t\t\n",
        "\t\t# While low's top two bits are 01 and high's are 10, delete the second highest bit of both\n",
        "\t\twhile (self.low & ~self.high & self.quarter_range) != 0:\n",
        "\t\t\tself.underflow()\n",
        "\t\t\tself.low = (self.low << 1) ^ self.half_range\n",
        "\t\t\tself.high = ((self.high ^ self.half_range) << 1) | self.half_range | 1\n",
        "\t\n",
        "\t\n",
        "\t# Called to handle the situation when the top bit of 'low' and 'high' are equal.\n",
        "\tdef shift(self):\n",
        "\t\traise NotImplementedError()\n",
        "\t\n",
        "\t\n",
        "\t# Called to handle the situation when low=01(...) and high=10(...).\n",
        "\tdef underflow(self):\n",
        "\t\traise NotImplementedError()\n",
        "\n",
        "\n",
        "# Encodes symbols and writes to an arithmetic-coded bit stream.\n",
        "class ArithmeticEncoder(ArithmeticCoderBase):\n",
        "\t\n",
        "\t# Constructs an arithmetic coding encoder based on the given bit output stream.\n",
        "\tdef __init__(self, numbits, bitout):\n",
        "\t\tsuper(ArithmeticEncoder, self).__init__(numbits)\n",
        "\t\t# The underlying bit output stream.\n",
        "\t\tself.output = bitout\n",
        "\t\t# Number of saved underflow bits. This value can grow without bound.\n",
        "\t\tself.num_underflow = 0\n",
        "\t\n",
        "\t\n",
        "\t# Encodes the given symbol based on the given frequency table.\n",
        "\t# This updates this arithmetic coder's state and may write out some bits.\n",
        "\tdef write(self, freqs, symbol):\n",
        "\t\tself.update(freqs, symbol)\n",
        "\t\n",
        "\t\n",
        "\t# Terminates the arithmetic coding by flushing any buffered bits, so that the output can be decoded properly.\n",
        "\t# It is important that this method must be called at the end of the each encoding process.\n",
        "\t# Note that this method merely writes data to the underlying output stream but does not close it.\n",
        "\tdef finish(self):\n",
        "\t\tself.output.write(1)\n",
        "\t\n",
        "\t\n",
        "\tdef shift(self):\n",
        "\t\tbit = self.low >> (self.num_state_bits - 1)\n",
        "\t\tself.output.write(bit)\n",
        "\t\t\n",
        "\t\t# Write out the saved underflow bits\n",
        "\t\tfor _ in range(self.num_underflow):\n",
        "\t\t\tself.output.write(bit ^ 1)\n",
        "\t\tself.num_underflow = 0\n",
        "\t\n",
        "\t\n",
        "\tdef underflow(self):\n",
        "\t\tself.num_underflow += 1\n",
        "\n",
        "\n",
        "# Reads from an arithmetic-coded bit stream and decodes symbols.\n",
        "class ArithmeticDecoder(ArithmeticCoderBase):\n",
        "\t\n",
        "\t# Constructs an arithmetic coding decoder based on the\n",
        "\t# given bit input stream, and fills the code bits.\n",
        "\tdef __init__(self, numbits, bitin):\n",
        "\t\tsuper(ArithmeticDecoder, self).__init__(numbits)\n",
        "\t\t# The underlying bit input stream.\n",
        "\t\tself.input = bitin\n",
        "\t\t# The current raw code bits being buffered, which is always in the range [low, high].\n",
        "\t\tself.code = 0\n",
        "\t\tfor _ in range(self.num_state_bits):\n",
        "\t\t\tself.code = self.code << 1 | self.read_code_bit()\n",
        "\t\n",
        "\t\n",
        "\t# Decodes the next symbol based on the given frequency table and returns it.\n",
        "\t# Also updates this arithmetic coder's state and may read in some bits.\n",
        "\tdef read(self, freqs):\n",
        "\t\t#if not isinstance(freqs, CheckedFrequencyTable):\n",
        "\t\t#\tfreqs = CheckedFrequencyTable(freqs)\n",
        "\t\t\n",
        "\t\t# Translate from coding range scale to frequency table scale\n",
        "\t\ttotal = int(freqs[-1])\n",
        "\t\t#total = freqs.get_total()\n",
        "\t\t#if total > self.maximum_total:\n",
        "\t\t#\traise ValueError(\"Cannot decode symbol because total is too large\")\n",
        "\t\trange = self.high - self.low + 1\n",
        "\t\toffset = self.code - self.low\n",
        "\t\tvalue = ((offset + 1) * total - 1) // range\n",
        "\t\t#assert value * range // total <= offset\n",
        "\t\t#assert 0 <= value < total\n",
        "\t\t\n",
        "\t\t# A kind of binary search. Find highest symbol such that freqs.get_low(symbol) <= value.\n",
        "\t\tstart = 0\n",
        "\t\tend = len(freqs)\n",
        "\t\t#end = freqs.get_symbol_limit()\n",
        "\t\twhile end - start > 1:\n",
        "\t\t\tmiddle = (start + end) >> 1\n",
        "\t\t\tlow = int(freqs[middle-1]) if middle > 0 else 0\n",
        "\t\t\t#if freqs.get_low(middle) > value:\n",
        "\t\t\tif low > value:\n",
        "\t\t\t\tend = middle\n",
        "\t\t\telse:\n",
        "\t\t\t\tstart = middle\n",
        "\t\t#assert start + 1 == end\n",
        "\t\t\n",
        "\t\tsymbol = start\n",
        "\t\t#assert freqs.get_low(symbol) * range // total <= offset < freqs.get_high(symbol) * range // total\n",
        "\t\tself.update(freqs, symbol)\n",
        "\t\t#if not (self.low <= self.code <= self.high):\n",
        "\t\t#\traise AssertionError(\"Code out of range\")\n",
        "\t\treturn symbol\n",
        "\t\n",
        "\t\n",
        "\tdef shift(self):\n",
        "\t\tself.code = ((self.code << 1) & self.state_mask) | self.read_code_bit()\n",
        "\t\n",
        "\t\n",
        "\tdef underflow(self):\n",
        "\t\tself.code = (self.code & self.half_range) | ((self.code << 1) & (self.state_mask >> 1)) | self.read_code_bit()\n",
        "\t\n",
        "\t\n",
        "\t# Returns the next bit (0 or 1) from the input stream. The end\n",
        "\t# of stream is treated as an infinite number of trailing zeros.\n",
        "\tdef read_code_bit(self):\n",
        "\t\ttemp = self.input.read()\n",
        "\t\tif temp == -1:\n",
        "\t\t\ttemp = 0\n",
        "\t\treturn temp\n",
        "\n",
        "\n",
        "# ---- Bit-oriented I/O streams ----\n",
        "\n",
        "# A stream of bits that can be read. Because they come from an underlying byte stream,\n",
        "# the total number of bits is always a multiple of 8. The bits are read in big endian.\n",
        "class BitInputStream(object):\n",
        "\t\n",
        "\t# Constructs a bit input stream based on the given byte input stream.\n",
        "\tdef __init__(self, inp):\n",
        "\t\t# The underlying byte stream to read from\n",
        "\t\tself.input = inp\n",
        "\t\t# Either in the range [0x00, 0xFF] if bits are available, or -1 if end of stream is reached\n",
        "\t\tself.currentbyte = 0\n",
        "\t\t# Number of remaining bits in the current byte, always between 0 and 7 (inclusive)\n",
        "\t\tself.numbitsremaining = 0\n",
        "\t\n",
        "\t\n",
        "\t# Reads a bit from this stream. Returns 0 or 1 if a bit is available, or -1 if\n",
        "\t# the end of stream is reached. The end of stream always occurs on a byte boundary.\n",
        "\tdef read(self):\n",
        "\t\tif self.currentbyte == -1:\n",
        "\t\t\treturn -1\n",
        "\t\tif self.numbitsremaining == 0:\n",
        "\t\t\ttemp = self.input.read(1)\n",
        "\t\t\tif len(temp) == 0:\n",
        "\t\t\t\tself.currentbyte = -1\n",
        "\t\t\t\treturn -1\n",
        "\t\t\tself.currentbyte = temp[0] if python3 else ord(temp)\n",
        "\t\t\tself.numbitsremaining = 8\n",
        "\t\tassert self.numbitsremaining > 0\n",
        "\t\tself.numbitsremaining -= 1\n",
        "\t\treturn (self.currentbyte >> self.numbitsremaining) & 1\n",
        "\t\n",
        "\t\n",
        "\t# Reads a bit from this stream. Returns 0 or 1 if a bit is available, or raises an EOFError\n",
        "\t# if the end of stream is reached. The end of stream always occurs on a byte boundary.\n",
        "\tdef read_no_eof(self):\n",
        "\t\tresult = self.read()\n",
        "\t\tif result != -1:\n",
        "\t\t\treturn result\n",
        "\t\telse:\n",
        "\t\t\traise EOFError()\n",
        "\t\n",
        "\t\n",
        "\t# Closes this stream and the underlying input stream.\n",
        "\tdef close(self):\n",
        "\t\tself.input.close()\n",
        "\t\tself.currentbyte = -1\n",
        "\t\tself.numbitsremaining = 0\n",
        "\n",
        "\n",
        "# A stream where bits can be written to. Because they are written to an underlying\n",
        "# byte stream, the end of the stream is padded with 0's up to a multiple of 8 bits.\n",
        "# The bits are written in big endian.\n",
        "class BitOutputStream(object):\n",
        "\t\n",
        "\t# Constructs a bit output stream based on the given byte output stream.\n",
        "\tdef __init__(self, out):\n",
        "\t\tself.output = out  # The underlying byte stream to write to\n",
        "\t\tself.currentbyte = 0  # The accumulated bits for the current byte, always in the range [0x00, 0xFF]\n",
        "\t\tself.numbitsfilled = 0  # Number of accumulated bits in the current byte, always between 0 and 7 (inclusive)\n",
        "\t\n",
        "\t\n",
        "\t# Writes a bit to the stream. The given bit must be 0 or 1.\n",
        "\tdef write(self, b):\n",
        "\t\tif b not in (0, 1):\n",
        "\t\t\traise ValueError(\"Argument must be 0 or 1\")\n",
        "\t\tself.currentbyte = (self.currentbyte << 1) | b\n",
        "\t\tself.numbitsfilled += 1\n",
        "\t\tif self.numbitsfilled == 8:\n",
        "\t\t\ttowrite = bytes((self.currentbyte,)) if python3 else chr(self.currentbyte)\n",
        "\t\t\tself.output.write(towrite)\n",
        "\t\t\tself.currentbyte = 0\n",
        "\t\t\tself.numbitsfilled = 0\n",
        "\t\n",
        "\t\n",
        "\t# Closes this stream and the underlying output stream. If called when this\n",
        "\t# bit stream is not at a byte boundary, then the minimum number of \"0\" bits\n",
        "\t# (between 0 and 7 of them) are written as padding to reach the next byte boundary.\n",
        "\tdef close(self):\n",
        "\t\twhile self.numbitsfilled != 0:\n",
        "\t\t\tself.write(0)\n",
        "\t\tself.output.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mjEj_H_Pz1vr"
      },
      "source": [
        "## Compress"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4jZ5Dga8hMsO",
        "cellView": "form"
      },
      "source": [
        "#@title Preprocess\n",
        "\n",
        "if mode != 'decompress':\n",
        "  input_path = path_to_file\n",
        "\n",
        "  if preprocess == 'cmix':\n",
        "    !./cmix/cmix -s ./cmix/dictionary/english.dic $path_to_file ./data/preprocessed.dat\n",
        "    input_path = \"./data/preprocessed.dat\"\n",
        "\n",
        "  # int_list will contain the characters of the file.\n",
        "  int_list = []\n",
        "  if preprocess == 'nncp' or preprocess == 'nncp-done':\n",
        "    if preprocess == 'nncp':\n",
        "      !time ./nncp-2019-11-16/preprocess c data/dictionary.words $path_to_file data/preprocessed.dat $n_words $min_freq\n",
        "    else:\n",
        "      !cp $path_to_file data/preprocessed.dat\n",
        "    input_path = \"./data/preprocessed.dat\"\n",
        "    orig = open(input_path, 'rb').read()\n",
        "    for i in range(0, len(orig), 2):\n",
        "      int_list.append(orig[i] * 256 + orig[i+1])\n",
        "    vocab_size = int(subprocess.check_output(\n",
        "        ['wc', '-l', 'data/dictionary.words']).split()[0])\n",
        "  else:\n",
        "    text = open(input_path, 'rb').read()\n",
        "    vocab = sorted(set(text))\n",
        "    vocab_size = len(vocab)\n",
        "    # Creating a mapping from unique characters to indexes.\n",
        "    char2idx = {u:i for i, u in enumerate(vocab)}\n",
        "    for idx, c in enumerate(text):\n",
        "      int_list.append(char2idx[c])\n",
        "\n",
        "  # Round up to a multiple of 8 to improve performance.\n",
        "  vocab_size = math.ceil(vocab_size/8) * 8\n",
        "  file_len = len(int_list)\n",
        "  print ('Length of file: {} symbols'.format(file_len))\n",
        "  print ('Vocabulary size: {}'.format(vocab_size))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VQ_3o0ZeEKcR",
        "cellView": "form"
      },
      "source": [
        "#@title Compression\n",
        "\n",
        "if mode == 'compress' or mode == 'both':\n",
        "  original_file = path_to_file\n",
        "  path_to_file = \"data/compressed.dat\"\n",
        "  with open(path_to_file, \"wb\") as out, contextlib.closing(BitOutputStream(out)) as bitout:\n",
        "    length = len(int_list)\n",
        "    # Write the original file length to the compressed file header.\n",
        "    out.write(length.to_bytes(5, byteorder='big', signed=False))\n",
        "    if preprocess != 'nncp' and preprocess != 'nncp-done':\n",
        "      # If NNCP was not used for preprocessing, write 256 bits to the compressed\n",
        "      # file header to keep track of the vocabulary.\n",
        "      for i in range(256):\n",
        "        if i in char2idx:\n",
        "          bitout.write(1)\n",
        "        else:\n",
        "          bitout.write(0)\n",
        "    enc = ArithmeticEncoder(32, bitout)\n",
        "    process(True, length, vocab_size, enc, int_list)\n",
        "  print(\"Compressed size:\", os.path.getsize(path_to_file))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YEmm5PznEk5q",
        "cellView": "form"
      },
      "source": [
        "#@title Download Result\n",
        "\n",
        "if mode == 'preprocess_only':\n",
        "  if preprocess == 'nncp':\n",
        "    download('data/dictionary.words')\n",
        "  download(input_path)\n",
        "elif mode != 'decompress':\n",
        "  download('data/compressed.dat')\n",
        "  if preprocess == 'nncp':\n",
        "    download('data/dictionary.words')\n",
        "  if checkpoint and mode != \"both\":\n",
        "    download('data/model.index')\n",
        "    download('data/model.data-00000-of-00001')\n",
        "    download('data/checkpoint')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H2hIA7_R5F2m"
      },
      "source": [
        "## Decompress"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hZkjku3phchb",
        "cellView": "form"
      },
      "source": [
        "#@title Decompression\n",
        "\n",
        "if mode == 'decompress' or mode == 'both':\n",
        "  output_path = \"data/decompressed.dat\"\n",
        "  with open(path_to_file, \"rb\") as inp, open(output_path, \"wb\") as out:\n",
        "    # Read the original file size from the header.\n",
        "    length = int.from_bytes(inp.read()[:5], byteorder='big')\n",
        "    inp.seek(5)\n",
        "    # Create a list to store the file characters.\n",
        "    output = [0] * length\n",
        "    bitin = BitInputStream(inp)\n",
        "    if preprocess == 'nncp' or preprocess == 'nncp-done':\n",
        "      # If the preprocessor is NNCP, we can get the vocab_size from the\n",
        "      # dictionary.\n",
        "      vocab_size = int(subprocess.check_output(\n",
        "          ['wc', '-l', 'data/dictionary.words']).split()[0])\n",
        "    else:\n",
        "      # If the preprocessor is not NNCP, we can get the vocabulary from the file\n",
        "      # header.\n",
        "      vocab = []\n",
        "      for i in range(256):\n",
        "        if bitin.read():\n",
        "          vocab.append(i)\n",
        "      vocab_size = len(vocab)\n",
        "    # Round up to a multiple of 8 to improve performance.\n",
        "    vocab_size = math.ceil(vocab_size/8) * 8\n",
        "    dec = ArithmeticDecoder(32, bitin)\n",
        "    process(False, length, vocab_size, dec, output)\n",
        "    # The decompressed data is stored in the \"output\" list. We can now write the\n",
        "    # data to file (based on the type of preprocessing used).\n",
        "    if preprocess == 'nncp' or preprocess == 'nncp-done':\n",
        "      for i in range(length):\n",
        "        out.write(bytes(((output[i] // 256),)))\n",
        "        out.write(bytes(((output[i] % 256),)))\n",
        "    else:\n",
        "      # Convert indexes back to the original characters.\n",
        "      idx2char = np.array(vocab)\n",
        "      for i in range(length):\n",
        "        out.write(bytes((idx2char[output[i]],)))\n",
        "\n",
        "  if preprocess == 'cmix':\n",
        "    !./cmix/cmix -d ./cmix/dictionary/english.dic $output_path ./data/final.dat\n",
        "    output_path = \"data/final.dat\"\n",
        "  if preprocess == 'nncp' or preprocess == 'nncp-done':\n",
        "    !./nncp-2019-11-16/preprocess d data/dictionary.words $output_path ./data/final.dat\n",
        "    output_path = \"data/final.dat\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ndF-mn69IJWF",
        "cellView": "form"
      },
      "source": [
        "#@title Download Result\n",
        "\n",
        "if mode == 'decompress':\n",
        "  if preprocess == 'nncp-done':\n",
        "    download('data/decompressed.dat')\n",
        "  else:\n",
        "    download(output_path)\n",
        "  if checkpoint:\n",
        "    download('data/model.index')\n",
        "    download('data/model.data-00000-of-00001')\n",
        "    download('data/checkpoint')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A5oyae1d9f0X",
        "cellView": "form"
      },
      "source": [
        "#@title Validation\n",
        "\n",
        "if mode == 'decompress' or mode == 'both':\n",
        "  if preprocess == 'nncp-done':\n",
        "    !md5sum data/decompressed.dat\n",
        "  !md5sum $output_path\n",
        "if mode == 'both':\n",
        "  !md5sum $original_file"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}